import streamlit as st
import os
import sys
import time
import logging
from reg_retrieval import ragdemo
from ollamaclienttest import OllamaClient

# 配置日志系统
# 使用绝对路径确保日志文件创建在正确位置
log_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'web_console.log')
logging.basicConfig(
    filename=log_file_path,
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s',
    encoding='utf-8'
)
print(f"[调试] 日志文件路径: {log_file_path}")
logger = logging.getLogger(__name__)

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 设置页面配置
st.set_page_config(
    page_title="文档问答系统",
    page_icon=":books:",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 初始化会话状态
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

if 'rag_instance' not in st.session_state:
    try:
        st.session_state.rag_instance = ragdemo()
        st.sidebar.success("向量数据库和模型初始化成功")
    except Exception as e:
        st.sidebar.error(f"初始化失败: {str(e)}")
        st.session_state.rag_instance = None

# 侧边栏配置
st.sidebar.title("模型配置")

# 模型来源选择
model_source = st.sidebar.radio(
    '模型来源',
    ['本地模型文件', 'Ollama服务'],
    index=1  # 默认选择Ollama服务
)
logger.debug(f"用户选择的模型来源: {model_source}")
print(f"[调试] 用户选择的模型来源: {model_source}")

# 根据选择显示不同的配置选项
if model_source == '本地模型文件':
    st.sidebar.info('请注意：使用本地模型需要确保模型目录包含必要的模型文件')
else:
    st.sidebar.info('使用Ollama服务不需要本地模型文件，但需要确保Ollama服务已安装并运行')

# 设置模型类型变量以便后续代码使用
model_type = model_source

# 根据模型类型选择模型
model_path = None
if model_type == "本地模型文件":
    # 获取用户的模型目录
    # 使用用户提供的模型目录作为默认值
    default_model_dir = os.path.join("C:\\", "ai", "models", "Imstudio-community")
    # 确保路径格式正确
    model_dir = st.sidebar.text_input("模型目录", default_model_dir, key="model_dir_input")
    model_dir = os.path.normpath(model_dir)  # 标准化路径格式
    logger.debug(f"模型目录: {model_dir}")
    print(f"[调试] 模型目录: {model_dir}")
    
    # 检查目录是否存在
    if os.path.exists(model_dir) and os.path.isdir(model_dir):
        logger.debug(f"模型目录存在: {model_dir}")
        print(f"[调试] 模型目录存在: {model_dir}")
        # 获取目录下的所有子目录并排序
        model_folders = sorted([f for f in os.listdir(model_dir) if os.path.isdir(os.path.join(model_dir, f))])
        logger.debug(f"找到的模型文件夹数量: {len(model_folders)}")
        print(f"[调试] 找到的模型文件夹数量: {len(model_folders)}")
        if model_folders:
            # 让用户选择模型
            selected_model = st.sidebar.selectbox("选择本地模型", model_folders, key="selected_model_selectbox")
            model_path = os.path.normpath(os.path.join(model_dir, selected_model))
            logger.debug(f"用户选择的本地模型路径: {model_path}")
            print(f"[调试] 用户选择的本地模型路径: {model_path}")
        else:
            st.sidebar.warning("模型目录为空，请放入模型文件夹后重试。\n您可以从Imstudio-community仓库下载模型文件到该目录。\n注意：模型文件夹中需要包含pytorch_model.bin、model.safetensors等模型文件。")
            logger.debug("模型目录为空")
            print("[调试] 模型目录为空")
            # 提供一个示例模型路径
            st.sidebar.info(f"示例模型路径: {os.path.join(model_dir, 'qwen3-7b')}")
            st.sidebar.info("模型文件应包含: pytorch_model.bin, model.safetensors, tf_model.h5, model.ckpt.index 或 flax_model.msgpack 之一")
    else:
        st.sidebar.error(f"模型目录不存在或不是有效目录: {model_dir}\n请创建该目录或修改路径")
        logger.debug(f"模型目录不存在或不是有效目录: {model_dir}")
        print(f"[调试] 模型目录不存在或不是有效目录: {model_dir}")
        # 提供一个按钮来创建目录
        if st.sidebar.button("创建模型目录"):
            try:
                os.makedirs(model_dir, exist_ok=True)
                st.sidebar.success(f"成功创建模型目录: {model_dir}")
                logger.debug(f"成功创建模型目录: {model_dir}")
                print(f"[调试] 成功创建模型目录: {model_dir}")
                # 刷新页面以更新目录状态
                st.experimental_rerun()
            except Exception as e:
                st.sidebar.error(f"创建目录失败: {str(e)}")
                logger.error(f"创建目录失败: {str(e)}")
                print(f"[调试] 错误: 创建目录失败: {str(e)}")
else:
    logger.debug("用户选择了Ollama服务模型")
    print("[调试] 用户选择了Ollama服务模型")

# 检查并初始化ollama客户端
# 强制重新初始化客户端，确保使用最新的模型配置
logger.debug(f"当前模型类型: {model_type}")
print(f"[调试] 当前模型类型: {model_type}")
logger.debug(f"当前模型路径: {model_path}")
print(f"[调试] 当前模型路径: {model_path}")
try:
    if model_type == "本地模型文件":
        if model_path:
            logger.debug(f"初始化本地模型客户端，路径: {model_path}")
            print(f"[调试] 初始化本地模型客户端，路径: {model_path}")
            # 确保传递绝对路径
            abs_model_path = os.path.abspath(model_path)
            logger.debug(f"转换为绝对路径: {abs_model_path}")
            print(f"[调试] 转换为绝对路径: {abs_model_path}")
            # 添加更多调试信息，检查路径是否存在
            if os.path.exists(abs_model_path):
                logger.debug(f"绝对路径存在: {abs_model_path}")
                print(f"[调试] 绝对路径存在: {abs_model_path}")
                # 列出模型目录下的文件，帮助调试
                try:
                    model_files = os.listdir(abs_model_path)
                    logger.debug(f"模型目录内容: {model_files}")
                    print(f"[调试] 模型目录内容: {model_files}")
                except Exception as e:
                    logger.error(f"无法列出模型目录内容: {str(e)}")
                    print(f"[调试] 错误: 无法列出模型目录内容: {str(e)}")
                
                try:
                    st.session_state.ollama_client = OllamaClient(model_path=abs_model_path)
                    st.sidebar.success(f"成功加载本地模型: {os.path.basename(abs_model_path)}")
                    # 测试模型是否可用
                    try:
                        test_response = st.session_state.ollama_client.generate(
                            model_name=os.path.basename(abs_model_path),
                            prompt="你好",
                            options={'temperature': 0.4}
                        )
                        logger.debug(f"模型测试成功，响应: {test_response[:20]}...")
                        print(f"[调试] 模型测试成功，响应: {test_response[:20]}...")
                    except Exception as e:
                        logger.error(f"模型测试失败: {str(e)}")
                        print(f"[调试] 错误: 模型测试失败: {str(e)}")
                        st.sidebar.warning(f"模型加载成功，但测试失败: {str(e)}")
                except Exception as e:
                    logger.error(f"初始化Ollama客户端失败: {str(e)}")
                    print(f"[调试] 错误: 初始化Ollama客户端失败: {str(e)}")
                    st.sidebar.error(f"加载模型失败: {str(e)}")
                    st.session_state.ollama_client = None
            else:
                logger.error(f"绝对路径不存在: {abs_model_path}")
                print(f"[调试] 错误: 绝对路径不存在: {abs_model_path}")
                st.sidebar.error(f"模型路径不存在: {abs_model_path}")
                st.session_state.ollama_client = None
        else:
            logger.debug("模型路径无效，无法初始化本地模型客户端")
            print("[调试] 模型路径无效，无法初始化本地模型客户端")
            st.sidebar.error("模型路径无效，请检查模型目录设置")
            st.session_state.ollama_client = None
    else:
        logger.debug("初始化Ollama服务客户端")
        print("[调试] 初始化Ollama服务客户端")
        # 添加Ollama服务端口配置
        ollama_port = st.sidebar.text_input("Ollama服务端口", "11434", key="ollama_port_input")
        logger.debug(f"Ollama服务端口: {ollama_port}")
        print(f"[调试] Ollama服务端口: {ollama_port}")
        # 对于Ollama服务，不需要提供模型路径，但可以指定端口
        base_url = f"http://localhost:{ollama_port}"
        try:
            st.session_state.ollama_client = OllamaClient(base_url=base_url)
            st.sidebar.success(f"Ollama客户端初始化成功，连接到: {base_url}")
            # 测试连接
            try:
                import requests
                test_url = f"{base_url}/api/tags"
                response = requests.get(test_url, timeout=5)
                if response.status_code == 200:
                    st.sidebar.info("Ollama服务连接测试成功")
                    logger.debug("Ollama服务连接测试成功")
                    print("[调试] Ollama服务连接测试成功")
                else:
                    st.sidebar.warning(f"Ollama服务连接测试失败，状态码: {response.status_code}")
                    logger.warning(f"Ollama服务连接测试失败，状态码: {response.status_code}")
                    print(f"[调试] Ollama服务连接测试失败，状态码: {response.status_code}")
            except Exception as e:
                st.sidebar.warning(f"Ollama服务连接测试失败: {str(e)}")
                logger.warning(f"Ollama服务连接测试失败: {str(e)}")
                print(f"[调试] Ollama服务连接测试失败: {str(e)}")
        except Exception as e:
            logger.error(f"初始化Ollama客户端失败: {str(e)}")
            print(f"[调试] 错误: 初始化Ollama客户端失败: {str(e)}")
            st.sidebar.error(f"Ollama客户端初始化失败: {str(e)}")
            st.session_state.ollama_client = None
    # 保存当前模型类型和路径
    st.session_state.last_model_type = model_type
    st.session_state.last_model_path = model_path
except Exception as e:
    logger.error(f"客户端初始化异常: {str(e)}")
    print(f"客户端初始化异常: {str(e)}")
    st.sidebar.error(f"客户端初始化失败: {str(e)}")
    st.session_state.ollama_client = None

# 其他侧边栏配置

# 模型选择（如果是Ollama服务模型）
if model_type != "本地模型文件":
    # Ollama服务模型选择
    model_name = st.sidebar.selectbox(
            "选择Ollama模型",
            ["qwen3:14b", "qwen3-30b-a3b-2507-thinking:256k", "llama3", "gemma:2b", "phi3:3.8b"],
            index=0
        )
else:
    # 对于本地模型，我们已经在前面选择了模型路径
    # 这里可以显示当前选择的模型名称
    if model_path:
        model_name = os.path.basename(model_path)
    else:
        model_name = "未选择有效模型"

temperature = st.sidebar.slider("温度参数", 0.0, 1.0, 0.4, 0.1)

# 主页面标题
st.title("文档问答系统")

# 显示聊天历史
for i, (question, answer, references) in enumerate(st.session_state.chat_history):
    with st.expander(f"问题 {i+1}: {question}"):
        st.write("**答案:**")
        st.write(answer)
        
        st.write("**参考来源:**")
        for j, ref in enumerate(references):
            st.write(f"{j+1}. 相似度: {ref['similarity_score']}")
            st.text_area(f"内容 {j+1}", ref['content'], height=100)

# 用户输入
user_question = st.text_input("请输入您的问题:", key="user_question")

# 提交按钮
if st.button("提交") and user_question and st.session_state.rag_instance and st.session_state.ollama_client:
    with st.spinner("正在检索知识库并生成答案..."):
        # 调用RAG检索
        rag_result = st.session_state.rag_instance.query(user_question)
        
        # 如果RAG没有找到相关信息，直接使用本地模型
        if rag_result['answer'] == "未在知识库中找到相关信息":
            context = "没有找到相关上下文信息"
        else:
            context = rag_result['context']
        
        # 构造提示词
        prompt = f"""根据下面的上下文内容回答问题：
[上下文]
{context}
[问题]
{user_question}
[要求]
- 如果上下文不足以回答问题，请回复‘根据现有资料无法确认’,
- 不要编造信息，
- 使用中文回答
"""
        
        # 调用模型
        answer = st.session_state.ollama_client.generate(
            model_name=model_name,
            prompt=prompt,
            options={'temperature': temperature}
        )
        
        # 保存到聊天历史
        st.session_state.chat_history.append((user_question, answer, rag_result['references']))
        
        # 显示结果
        st.write("**答案:**")
        st.write(answer)
        
        st.write("**参考来源:**")
        for j, ref in enumerate(rag_result['references']):
            st.write(f"{j+1}. 相似度: {ref['similarity_score']}")
            st.text_area(f"内容 {j+1}", ref['content'], height=100)

# 清除历史按钮
if st.button("清除历史"):
    st.session_state.chat_history = []
    st.experimental_rerun()

# 底部信息
st.markdown("---")
st.markdown("文档问答系统 v1.0 | 基于本地模型和向量数据库")